import torch
import torch_npu
from torch_geometric.loader import DataLoader
from tqdm import tqdm
from dc_error import eval_loss_fn
from networks.MPN import MPN, MPN_simplenet, SkipMPN, MaskEmbdMPN, MultiConvNet, MultiMPN, MaskEmbdMultiMPN
import argparse
from utils.evaluation import evaluate_epoch
from utils.custom_loss_functions import Masked_L2_loss


def argument_parser():
    parser = argparse.ArgumentParser(
        description='powerflownet inference'
    )
    # Network Parameters
    parser.add_argument('--model_path', type=str, required=True, help='Path of the trained model')
    parser.add_argument('--data_path', type=str, default='./data', help='Data directory path')
    parser.add_argument('--case', type=str, default='118v2', help='Grid case (e.g. 118v2,118)')
    parser.add_argument('--model', type=str, default='MaskEmbdMultiMPN', help='Model name')
    parser.add_argument('--batch_size', type=int, default=128, help='Batch size for inference')
    parser.add_argument('--disable_normalize', type=int, default=128, help='Batch size for inference')
    return parser.parse_args()
def main():
    args = argument_parser()
    model_path = args.model_path
    data_path = args.data_path
    grid_case = args.case
    batch_size = args.batch_size
    nomalize_data = not args.disable_normalize

    # 设置 device
    device = torch.device('cuda' if torch_npu.is_available() else 'cpu')
    print(f"Using device: {device}")
    # 加载数据
    testset = PowerFlowData(root=data_path, case=grid_case, split=[.5, .2, .3], task='test',
                            normalize=nomalize_data)
    test_loader = DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=4,prefetch_factor=4,pin_memory=True,perfomance_workers=True)
    # 加载权重
    checkpoint = torch.load(model_path, map_location=device)
    # 模型映射
    models = {
        'MPN': MPN,
        'MPN_simplenet': MPN_simplenet,
        'SkipMPN': SkipMPN,
        'MaskEmbdMPN': MaskEmbdMPN,
        'MultiConvNet': MultiConvNet,
        'MultiMPN': MultiMPN,
        'MaskEmbdMultiMPN': MaskEmbdMultiMPN
    }
    # 实例化模型
    model_class = models[args.model]
    node_in_dim, node_out_dim, edge_dim = testset.get_data_dimensions()

    model = model_class(
        nfeature_dim=node_in_dim,
        efeature_dim=edge_dim,
        output_dim=node_out_dim,
        hidden_dim=checkpoint['args'].hidden_dim,
        n_gnn_layers=checkpoint['args'].n_gnn_layers,
        K=checkpoint['args'].K,
        dropout_rate=checkpoint['args'].dropout_rate,
    ).to(device)
    # 处理DDP模型权重
    state_dict = checkpoint['model_state_dict']
    if next(iter(state_dict.keys())).startswith('module.'):
        state_dict = {k[7:]:v for k, v in state_dict.items()}

    model.load_state_dict(state_dict)
    model = model.eval()

    print(f"Model loaded from: {model_path}")
    print(f"Trained until epoch: {checkpoint['epoch']}")

    # 推理
    print("\nStarting inference...")
    eval_loss_fn = Masked_L2_loss(regularize=False)
    test_loss = evaluate_epoch(model, test_loader, eval_loss_fn, device)
    # 输出推理结果
    print("\n" + "="*50)
    print("inference results")
    print(f"Model:{args.model}")
    print(f"Graph case:{grid_case}")
    print(f"Test dataset size :{len(testset)}")
    print(f"Test loss: {test_loss:0.4f}")